Packed Sort: speed up MoE routing argsort via bit-packed single-key sort - #3488
Open
bigeye wants to merge 2 commits into
Open
Packed Sort: speed up MoE routing argsort via bit-packed single-key sort#3488bigeye wants to merge 2 commits into
bigeye wants to merge 2 commits into
Conversation
bigeye
requested review from
a1yssan13,
bythew3i,
gxd3,
jrplatin,
kyuyeunk and
lk-chen
as code owners
August 31, 2026 05:51
bigeye
force-pushed
the
donghyun/moesort-packedsort
branch
from
August 31, 2026 06:00
d864e01 to
0f9ad9d
Compare
MoE routing sorts tokens by expert with a stable argsort, which XLA runs
as a two-operand (key, index) sort. When the key and the index fit
together in an int32, pack them into one value, `key << shift | index`,
and run a plain single-key unstable sort instead. The index in the low
bits breaks ties, so the result is identical to the stable argsort, but
XLA has no iota operand and no stable-sort comparator to deal with.
The argsort for gather drops from 106.6 us to 21.7 us (4.9x speedup),
and the validity sort from 7.2us to 5.5us (1.3x speedup).
Measured on Qwen/Qwen3.5-397B-A17B-FP8, TPU v7x-8, preset 8k1k, 640
prompts, repeat=5, The two sides ran with:
Baseline: without any change. no env overrides.
Packed Sort + FP8 AG: MOE_ALL_GATHER_ACTIVATION_DTYPE=fp8
SC_ALLREDUCE_ALLGATHER_OFFLOAD_MIN_BYTES=0
MoE AG and Sort/Permute runs in parallel, so we need to optimize both at
the same time. Read the e2e deltas as the two changes together, not as the
sort alone.
| Conc | tok/s | std% | TTFT (ms) | TPOT (ms) |
|---|---|---|---|---|
| 64 | 15908.5 -> 16114.74 (+1.30%) | 0.99 -> 0.98 | 2014.13 -> 1999.67 (-0.72%) | 31.305 -> 30.906 (-1.27%) |
| 128 | 22883.2 -> 23027.63 (+0.63%) | 0.53 -> 0.56 | 2242.35 -> 2114.58 (-5.70%) | 44.922 -> 44.642 (-0.62%) |
| 256 | 28007.1 -> 30381.17 (+8.48%) | 5.48 -> 0.87 | 3258.17 -> 3127.76 (-4.00%) | 65.219 -> 64.296 (-1.42%) |
| 512 | 30888.1 -> 31775.83 (+2.87%) | 0.34 -> 0.39 | 30172.15 -> 28997.39 (-3.89%) | 94.396 -> 90.833 (-3.77%) |
Raw per-repeat throughputs for Packed Sort + FP8 AG:
c64 [15860.9, 16083.5, 16238.5, 16247.4, 16143.4]
c128 [22798.9, 23088.1, 23098.6, 23052.7, 23099.9]
c256 [30270.5, 30141.1, 30672.1, 30660.8, 30161.4]
c512 [31665.6, 31892.7, 31779.7, 31902.3, 31638.9]
| Eval | Baseline | Packed Sort + FP8 AG |
|---|---|---|
| mmlu_pro (700 q) * | 82.8 | 82.8 |
| gpqa_diamond_cot_zeroshot (50 q, flexible-extract) | 88.0 +/- 1.6 | 92.0 +/- 1.6 |
| aime25 (30 q) | 73.3 +/- 8.6 | 77.5 +/- 4.2 |
Signed-off-by: Donghyun Cho <donghyun@google.com>
bigeye
force-pushed
the
donghyun/moesort-packedsort
branch
from
August 31, 2026 06:01
0f9ad9d to
95e5674
Compare
Collaborator
Author
|
@qizzzh @guowei-dev for reviews. |
Collaborator
|
We've been exploring similar packed-sort ideas on the torchtpu side — https://github.com/vllm-project/vllm-torchtpu/pull/649 does the same routing packing, plus more of the routing prep, including a counting sort for the inverse permutation. If this is headed for torchtpu as well, maybe we can build on that directly? |
Collaborator
Author
|
@guowei-dev Thanks for the heads up! Good to know that there's some similar efforts being proven. I actually have a branch (will have a separate pr for that) for counting sort for the second sort as well. I've been working on that and saw the good performance gain there. As they are different repos at the end of the day, I think we can just apply these techniques in both repo in that it won't hurt the migration or future plans. Tell me if you have other thoughts. |
guowei-dev
approved these changes
Aug 31, 2026
guowei-dev
left a comment
guowei-dev
left a comment
Collaborator
There was a problem hiding this comment.
LGTM per offline discussion. The changes itself are in good shape, thanks for working on this!
Signed-off-by: Donghyun Cho <donghyun@google.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
MoE routing sorts tokens by expert with a stable argsort, which XLA runs
as a two-operand (key, index) sort. When the key and the index fit
together in an int32, pack them into one value,
key << shift | index,and run a plain single-key unstable sort instead. The index in the low
bits breaks ties, so the result is identical to the stable argsort, but
XLA has no iota operand and no stable-sort comparator to deal with.
The argsort for gather drops from 106.6 us to 21.7 us (4.9x speedup),
and the validity sort from 7.2us to 5.5us (1.3x speedup).
Measured on Qwen/Qwen3.5-397B-A17B-FP8, TPU v7x-8, preset 8k1k, 640
prompts, repeat=5, The two sides ran with:
SC_ALLREDUCE_ALLGATHER_OFFLOAD_MIN_BYTES=0
MoE AG and Sort/Permute runs in parallel, so we need to optimize both at
the same time. Read the e2e deltas as the two changes together, not as the
sort alone.
Raw per-repeat throughputs for Packed Sort + FP8 AG:
c64 [15860.9, 16083.5, 16238.5, 16247.4, 16143.4]
c128 [22798.9, 23088.1, 23098.6, 23052.7, 23099.9]
c256 [30270.5, 30141.1, 30672.1, 30660.8, 30161.4]
c512 [31665.6, 31892.7, 31779.7, 31902.3, 31638.9]
Signed-off-by: Donghyun Cho donghyun@google.com
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure: